from nicegui import ui
import numpy as np
from random import random
import ctypes
from python_app.railsim.libapi import (
    ElectricLocomotive,
    ElectricMotorProperty,
    ElectricMotor,
    ELECTRIC_MOTOR_TYPE,
)

motorProperty = ElectricMotorProperty.create_value(
    "YJ85-A".encode("utf8"),
    ELECTRIC_MOTOR_TYPE.THREE_PHASED_ASYNC,
    1250,
    11852,
    2662,
    0.92,
    4.85,
)

# 应该一口气创建一系列的电动机
motor_array = ElectricMotor.create_multi(
    6,
).contents


print("ptr-array address", hex(ctypes.addressof(motor_array)))
loco_ptr, loco = ElectricLocomotive.create(5, 25, 25 * 6, 1.20, 6, motorProperty)


print("axleLoad", loco.base.axleLoad)


def update():
    chart_data = compute_chart_data()
    for k in echart.options.keys():
        echart.options[k] = chart_data[k]
    echart.update()


def modify_property(obj: object, attr: str, value):
    setattr(obj, attr, value)


def compute_chart_data():
    loco.update_parameters()
    max_traction_limit_at_start = max(
        loco.traction_limit_adhesion(0), loco.traction_limit_torque()
    )
    max_rpm_speed = loco.speed_limit_rpm()
    print(
        max_traction_limit_at_start,
        loco.traction_limit_adhesion(0),
        loco.traction_limit_torque(),
    )
    print(max_rpm_speed)
    CHART_XAXIS_MAX = max_rpm_speed * 1.05
    max_traction_at_max_speed = max(
        loco.traction_limit_power(loco.base.maxPower, max_rpm_speed),
        loco.traction_limit_adhesion(max_rpm_speed),
    )
    return {
        "tooltip": {
            "trigger": "axis",
            ":formatter": """function (params) {
                var relVal = params[0].name;
                for (var i = 0, l = params.length; i < l; i++) {
                    relVal += '<br/>' + params[i].marker + params[i].seriesName + ' : ' + Number(params[i].value[1]).toFixed(2)
                }
                return relVal;
            }
            """,
        },
        "xAxis": {
            "type": "value",
        },
        "yAxis": {"type": "value"},
        "legend": {"textStyle": {"color": "gray"}},
        "series": [
            {
                "type": "line",
                "name": "牵引包线——轮轨黏着",
                "showSymbol": False,
                "data": [
                    [speed, loco.traction_limit_adhesion(speed)]
                    for speed in np.arange(0, CHART_XAXIS_MAX, 0.5)
                ],
            },
            {
                "type": "line",
                "name": f"牵引包线—电机功率@{loco.base.maxPower}",
                "showSymbol": False,
                "data": [
                    [speed, loco.traction_limit_power(loco.base.maxPower, speed)]
                    for speed in np.arange(0, CHART_XAXIS_MAX, 0.5)
                    if loco.traction_limit_power(loco.base.maxPower, speed)
                    < (max_traction_limit_at_start * 1.1)
                ],
            },
            {
                "type": "line",
                "name": f"牵引包线—电机极限扭矩",
                "showSymbol": False,
                "data": [
                    [speed, loco.traction_limit_torque()]
                    for speed in np.arange(0, CHART_XAXIS_MAX, 0.5)
                    if (loco.traction_limit_torque())
                    < (loco.traction_limit_power(loco.base.maxPower, speed) * 1.2)
                ],
            },
            {
                "type": "line",
                "name": f"牵引包线—电机极限转速",
                "showSymbol": False,
                "data": [
                    [max_rpm_speed, 0],
                    [
                        max_rpm_speed,
                        loco.traction_limit_power(loco.base.maxPower, max_rpm_speed),
                    ],
                    [
                        max_rpm_speed,
                        loco.traction_limit_adhesion(max_rpm_speed),
                    ],
                    [
                        max_rpm_speed,
                        max_traction_at_max_speed * 1.05,
                    ],
                ],
            },
        ],
    }


print("axleload", loco.base.axleLoad)
with ui.row().style("width: 100vw; height:100vh; display: flex;"):
    with ui.column().style("height: 100%; max-width:30%"):
        ui.number(
            label=f"电机最大功率(kW)x{loco.motorNum}轴",
            value=loco.motorProperty.maxPower,
            format="%.2f",
            on_change=lambda evt: modify_property(
                loco.motorProperty, "maxPower", evt.value
            ),
        )
        ui.number(
            label="驱动轴重(t)",
            value=loco.base.axleLoad,
            format="%.2f",
            on_change=lambda evt: modify_property(loco.base, "axleLoad", evt.value),
        )
        ui.number(
            label="驱动齿轮比",
            value=loco.motorProperty.gearRatio,
            format="%.2f",
            on_change=lambda evt: modify_property(
                loco.motorProperty, "gearRatio", evt.value
            ),
        )

        ui.button("Update", on_click=update)
    with ui.column().style("height: 100%; flex-grow: 1"):
        echart = ui.echart(compute_chart_data()).style("height: 100%")


ui.run(title="牵引特性曲线计算器")
